package com.jeehentor.ai.controller;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.metadata.IPage;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.jeehentor.ai.config.AiClientManager;
import com.jeehentor.ai.core.AiClient;
import com.jeehentor.ai.dto.ChatMessage;
import com.jeehentor.ai.dto.ChatRequest;
import com.jeehentor.ai.entity.AiChatMessage;
import com.jeehentor.ai.entity.AiChatSession;
import com.jeehentor.ai.mapper.AiPlatformConfigMapper;
import com.jeehentor.ai.service.AiChatMessageService;
import com.jeehentor.ai.service.AiChatSessionService;
import com.jeehentor.common.vo.Result;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import com.jeehentor.ai.entity.AiPlatformConfig;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.io.PrintWriter;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @Description : TODO
 * @Author : wzd
 * @Date : 2025/4/17 14:28
 * @Version : 1.0
 **/
@RestController
@RequestMapping("/ai")
@Slf4j
public class AiChatController {

    @Autowired
    private AiClientManager aiClientManager;
    @Autowired
    private AiPlatformConfigMapper configMapper;
    @Autowired
    private AiChatMessageService messageService;
    @Autowired
    private AiChatSessionService sessionService;
    @GetMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public void streamChat(@RequestParam String platform,
                           @RequestParam String sessionId,
                           @RequestParam String messages,
                           HttpServletResponse response) {
        response.setContentType("text/event-stream;charset=UTF-8");

        AiClient client = aiClientManager.getClient(platform);

        List<ChatMessage> messageList;
        try {
            messageList = new ObjectMapper().readValue(messages, new TypeReference<List<ChatMessage>>() {});
        } catch (Exception e) {
            log.error("解析 messages 失败: {}", messages, e);
            response.setStatus(HttpServletResponse.SC_BAD_REQUEST);
            return;
        }

        Flux<String> stream = client.streamChat(messageList, sessionId);

        try (PrintWriter writer = response.getWriter()) {
            stream
                    .doOnNext(token -> {
                        log.info("token: {}", token);
                        writer.write("data: " + token + "\n\n");
                        writer.flush();
                    })
                    .doOnComplete(() -> {
                        writer.write("data: [DONE]\n\n");
                        writer.flush();
                    })
                    .blockLast();
        } catch (IOException e) {
            log.error("写入 SSE 响应失败", e);
        }
    }


    @PostMapping("/chat")
    public Result<String> chat(@RequestBody ChatRequest request) {
        String result = aiClientManager.getClient(request.getPlatform())
                .chat(request.getPrompt(), request.getSessionId());
        return Result.ok(result);
    }

    @GetMapping("/isStream")
    public Result<Boolean> isStream(@RequestParam String platform) {
        return Result.ok(aiClientManager.isStream(platform));
    }

    @GetMapping("/platforms")
    public Result<List<String>> listPlatforms() {
        return Result.ok(configMapper.selectList(null).stream()
                .filter(AiPlatformConfig::getIsActive)
                .map(AiPlatformConfig::getPlatform)
                .distinct()
                .collect(Collectors.toList()));
    }

    @GetMapping("/session/page")
    public Result<IPage<AiChatSession>> listPage(
            @RequestParam(defaultValue = "1") int page,
            @RequestParam(defaultValue = "20") int size) {
        Page<AiChatSession> p = new Page<>(page, size);
        LambdaQueryWrapper<AiChatSession> wrapper = Wrappers.lambdaQuery();
        wrapper.orderByDesc(AiChatSession::getCreateTime);
        return Result.ok(sessionService.page(p, wrapper));
    }

    @PostMapping("/session/save")
    public Result<Void> save(@RequestBody AiChatSession session) {
        LambdaQueryWrapper<AiChatSession> query = Wrappers.lambdaQuery();
        query.eq(AiChatSession::getSessionId, session.getSessionId());

        AiChatSession exist = sessionService.getOne(query, false);
        if (exist == null) {
            sessionService.save(session);
        }

        return Result.OK();
    }

    @GetMapping("/messages/page")
    public Result<IPage<AiChatMessage>> listMessage(
            @RequestParam String sessionId,
            @RequestParam(defaultValue = "1") int page,
            @RequestParam(defaultValue = "50") int size) {
        Page<AiChatMessage> p = new Page<>(page, size);
        LambdaQueryWrapper<AiChatMessage> wrapper = Wrappers.lambdaQuery();
        wrapper.eq(AiChatMessage::getSessionId, sessionId)
                .orderByAsc(AiChatMessage::getCreateTime);
        return Result.ok(messageService.page(p, wrapper));
    }

    @PostMapping("/message/save")
    public Result<Void> save(@RequestBody AiChatMessage message) {
        message.setCreateTime(new Date());
        messageService.save(message);
        return Result.OK();
    }

}